Capstone Project Group A CV1
Pneumonia Detection Challenge¶
The goal is to build a pneumonia detection system, to locate the position of inflammation in an image.
This model is built using Detectron2 (Faster_RCNN) on pytorch: https://github.com/facebookresearch/detectron2
Data preprocessing steps needed are as below:
- Convert all the dicom images to jpg as Detectron2 can't directly read dicom images : convert_to_jpg()
- Image augmentation using albumentations library
- Convert the data to coco json format : get_data_dicts()
In [1]:
!pip install pyyaml==5.1
!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.9/index.html
!pip install pydicom
!pip install sklearn
!pip install funcy
!pip install argparse
In [2]:
# check pytorch installation:
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
assert torch.__version__.startswith("1.9") # please manually install torch 1.9 if Colab changes its default version
In [3]:
# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
# import some common libraries
import numpy as np
import pandas as pd
import glob, os, json, cv2, random
from google.colab.patches import cv2_imshow
import pydicom as pyd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
import matplotlib.pyplot as plt
In [4]:
base_dir='/content/drive/MyDrive/colab'
data_dir = base_dir+'/input'
working_dir = base_dir+'/working'
train_labels_df = pd.read_csv(data_dir+'/stage_2_train_labels.csv')
label_meta_data = pd.read_csv(data_dir+'/stage_2_detailed_class_info.csv')
train_dicom_dir = os.path.join(data_dir, 'stage_2_train_images')
test_dicom_dir = os.path.join(data_dir, 'stage_2_test_images')
train_labels_df.tail(10)
Out[4]:
In [5]:
def get_files(dicom_dir):
files = glob.glob(dicom_dir+'/'+'*.dcm')
return list(set(files))
def parse_dataset(dicom_dir, anns):
image_files = get_files(dicom_dir)
image_annotations = {fp: [] for fp in image_files}
for index, row in anns.iterrows():
fp = os.path.join(dicom_dir, row['patientId']+'.dcm')
image_annotations[fp].append(row)
if index == 0:
print(row)
return image_files, image_annotations
image_files, image_annotations = parse_dataset(train_dicom_dir, anns=train_labels_df)
In [6]:
!pip install -q -U albumentations
!echo "$(pip freeze | grep albumentations) is successfully installed"
In [7]:
import albumentations as A
from albumentations.core.composition import BboxParams
transform = A.Compose([
A.RandomCrop(width=1024, height=1024),
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.2),
], bbox_params=BboxParams(format='coco', min_area=1024, min_visibility=0.1)
)
In [8]:
def convert_to_jpg(input_file, output_file):
ds = pyd.dcmread(input_file)
cv2.imwrite(output_file, ds.pixel_array)
# for d in ["test"]:
# input_path = os.path.join(data_dir, "stage_2_"+d+"_images")
# for file in get_files(input_path):
# output_file = os.path.join(data_dir, d, os.path.basename(file).replace(".dcm", "")+".jpg")
# convert_to_jpg(file, output_file)
In [9]:
def show(d):
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=pnuemonia_metadata, scale=0.5)
out = visualizer.draw_dataset_dict(d)
cv2_imshow(out.get_image()[:, :, ::-1])
In [10]:
from detectron2.structures import BoxMode
def get_data_dicts(img_dir, augment):
dataset_dicts = []
for idx, arr in enumerate(image_annotations.values()):
record = {}
filename = os.path.join(img_dir, arr[0]["patientId"]+".jpg")
img = cv2.imread(filename)
if img is None:
continue;
height, width = img.shape[:2]
record["file_name"] = filename
record["image_id"] = idx
record["height"] = height
record["width"] = width
objs = []
aug_boxes = []
for v in arr:
if np.isnan(float(v["x"])):
continue
obj = {
"bbox": [v["x"], v["y"], v["width"], v["height"]],
"bbox_mode": BoxMode.XYWH_ABS,
"category_id": 0,
}
aug_boxes.append([v["x"], v["y"], v["width"], v["height"], 'pnuemonia'])
objs.append(obj)
record["annotations"] = objs
# dataset_dicts.append(record)
rndm = random.randint(1, 10)
if augment and rndm < 3 and len(objs) > 0:
transformed = transform(image=img, bboxes=aug_boxes)
transformed_image = transformed['image']
transformed_bboxes = transformed['bboxes']
aug_file_name = os.path.join(img_dir,'augment', arr[0]["patientId"]+".jpg")
cv2.imwrite(aug_file_name, transformed_image)
# plt.imshow(cv2.cvtColor(transformed_image, cv2.COLOR_BGR2RGB))
aug_objs = []
for row in transformed_bboxes:
obj = {
"bbox": row[0:4],
"bbox_mode": BoxMode.XYWH_ABS,
"category_id": 0,
}
aug_objs.append(obj)
record_aug = {}
record_aug["image_id"] = idx
record_aug["height"] = height
record_aug["width"] = width
record_aug["annotations"] = aug_objs
record_aug["file_name"] = aug_file_name
record_aug["image_id"] = idx+40000
show(record)
show(record_aug)
print(filename)
print(aug_file_name)
break;
return dataset_dicts
In [12]:
pnuemonia_metadata = MetadataCatalog.get("train")
print( MetadataCatalog.get("train"))
get_data_dicts(data_dir+'/train', True)
Out[12]:
Create dataset without image augmentation¶
In [13]:
# import json
# out_file = open(os.path.join(data_dir,"coco-full.json"), "w")
# json.dump(get_data_dicts(data_dir+'/train', False), out_file, indent = 4)
# out_file.close()
Split the dataset into train and validation¶
In [73]:
with open(os.path.join(data_dir,"coco-full.json"), 'r') as myfile:
data=myfile.read()
# parse file
obj = json.loads(data)
X_train, X_val = train_test_split(obj, test_size=0.1)
print(len(X_train))
print(len(X_val))
In [15]:
def get_dicts(name):
return X_train if name == 'train' else X_val
for d in ["train", "val"]:
try:
DatasetCatalog.register(d, lambda d=d: get_dicts(d))
MetadataCatalog.get(d).set(thing_classes=["pneumonia"])
except:
print('')
In [24]:
from detectron2.engine import DefaultTrainer
global cfg
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_C4_1x.yaml"))
cfg.DATASETS.TRAIN = ("train")
cfg.DATASETS.TEST = ()
cfg.OUTPUT_DIR = ( base_dir + "/output")
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_C4_1x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 500
cfg.SOLVER.STEPS = []
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
def train_model():
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
def evaluate(cnt):
# Inference should use the config with parameters that are used in training
# cfg now already contains everything we've set previously. We changed it a little bit for inference:
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") # path to the model we just trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold
result = [];
ctr = 0;
for d in X_val:
im = cv2.imread(d["file_name"])
outputs = predictor(im)
no_objects = len(outputs["instances"].pred_boxes)
predicted = 1 if no_objects > 0 else 0
ground_truth = 1 if len(d["annotations"]) > 0 else 0
result.append({"id": os.path.basename(d["file_name"]).replace(".jpg", ""), "ground_truth": ground_truth, "predicted": predicted})
ctr +=1
if ctr > cnt:
break;
return result
In [74]:
train_model()
Inference & evaluation using the trained model¶
In [75]:
result = evaluate(3000)
In [76]:
df = pd.DataFrame(result)
df.head(10)
Out[76]:
In [77]:
print(classification_report(df['ground_truth'], df['predicted']))
accuracy = accuracy_score(df['ground_truth'], df['predicted'])
print('Model accuracy is: ', accuracy)
Create dataset with image augmentation for positive cases and save¶
In [29]:
# import json
# out_file = open(os.path.join(data_dir,"coco-full-aug.json"), "w")
# json.dump(get_data_dicts(data_dir+'/train', True), out_file, indent = 4)
# out_file.close()
Split the dataset into train and validation¶
In [30]:
with open(os.path.join(data_dir,"coco-full-aug.json"), 'r') as myfile:
data=myfile.read()
# parse file
obj = json.loads(data)
X_train, X_val = train_test_split(obj, test_size=0.1)
print(len(X_train))
print(len(X_val))
In [31]:
train_model()
Inference & evaluation using the trained model on augmented data¶
In [69]:
result = evaluate(3000)
In [70]:
df_aug = pd.DataFrame(result)
df_aug.head()
Out[70]:
In [71]:
print(classification_report(df_aug['ground_truth'], df_aug['predicted']))
accuracy = accuracy_score(df_aug['ground_truth'], df_aug['predicted'])
print('Model accuracy is: ', accuracy)
Make Predictions on Random Data¶
In [59]:
def get_test_files(dir):
files = glob.glob(dir+'/'+'*.jpg')
return list(set(files))
In [68]:
from detectron2.utils.visualizer import ColorMode
# get_test_files(data_dir+'/test')
for d in random.sample(X_val, 10):
im = cv2.imread(d['file_name'])
predictor = DefaultPredictor(cfg)
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=pnuemonia_metadata,
scale=0.5,
# instance_mode=ColorMode.IMAGE_BW
)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
print("\nPrediction")
cv2_imshow(out.get_image()[:, :, ::-1])
print("\nGround Truth")
show(d)
